package com.ycy.socket;

import cn.hutool.core.collection.CollUtil;
import cn.hutool.core.lang.mutable.MutableObj;

import java.util.*;
import java.util.concurrent.ConcurrentHashMap;

/**
 *
 * <p>
 *     这是一个提供了二级分桶功能的Map，本类会通过拓展KEY来对已有的Map结构进行再分桶。
 * </p>
 * <p>使用者可以通过原生KEY直接获取到对象，也可以通过拓展KEY获取整个分桶里的所有对象。</p>
 * <p>注意，分桶并不代表原生KEY之间相互隔离。所以使用者依然需要考虑原生KEY重复时，会导致的value被覆盖的问题。</p>
 *
 *
 * @param <RAW_KEY> 原生KEY类型
 * @param <EXT_KEY> 拓展KEY类型
 * @param <T> 实际存放的value类型
 *
 * @author PanBingYu
 */


public class ExtendKeyBucketMap<RAW_KEY, EXT_KEY, T> {

    private Map<EXT_KEY, ConcurrentHashMap<RAW_KEY, MutableObj<T>>> extMap = new ConcurrentHashMap<>();

    private Map<RAW_KEY, MutableObj<T>> rawMap = new ConcurrentHashMap<>();

    public ExtendKeyBucketMap() {
        this(true);
    }

    public ExtendKeyBucketMap(boolean isConcurrent) {
        if (!isConcurrent) {
            extMap = new HashMap<>();
            rawMap = new HashMap<>();
        }
    }

    public ExtendKeyBucketMap(Map<EXT_KEY, ConcurrentHashMap<RAW_KEY, MutableObj<T>>> extMap, Map<RAW_KEY, MutableObj<T>> rawMap) {
        this.extMap = extMap;
        this.rawMap = rawMap;
    }



    /**
     * 获取原生的二级拓展key的Map
     * @return
     */
    public Map<EXT_KEY, ConcurrentHashMap<RAW_KEY, MutableObj<T>>> getExtMap() {
        return extMap;
    }

    /**
     * 获取原生的全量Map
     * @return
     */
    public Map<RAW_KEY, MutableObj<T>> getRawKeyRawMap() {
        return rawMap;
    }

    /**
     * 获取所有原生Key
     * @return
     */
    public Set<RAW_KEY> rawKeySet() {
        return this.rawMap.keySet();
    }

    /**
     * 获取所有拓展key
     * @return
     */
    public Set<EXT_KEY> extKeySet() {
        return this.extMap.keySet();
    }


    /**
     * 指定一级原生key与二级拓展key，将对象放入map中
     * @param key
     * @param corpCode
     * @param value
     * @return
     */
    public T put(RAW_KEY key, EXT_KEY corpCode, T value) {
        ConcurrentHashMap<RAW_KEY, MutableObj<T>> container = extMap.get(corpCode);
        if (null == container) {
            container = new ConcurrentHashMap<>();
            extMap.put(corpCode, container);
        }
        MutableObj<T> wrapper = new MutableObj<>(value);

        // 放入corpCode拓展map
        container.put(key, wrapper);

        // 放入全量map
        rawMap.put(key, wrapper);

        return value;
    }

    /**
     * 根据一级原生key获取对象
     * @param key
     * @return
     */
    public T get(RAW_KEY key) {
        return Optional.ofNullable(rawMap.get(key))
                .map(MutableObj::get)
                .orElse(null);
    }

    /**
     * 指定一级原生key与二级拓展key，获取对象
     * @param key
     * @param corpCode
     * @return
     */
    public T get(RAW_KEY key, EXT_KEY corpCode) {
        ConcurrentHashMap<RAW_KEY, MutableObj<T>> container = extMap.get(corpCode);
        if (null == container || container.isEmpty()) {
            return null;
        }

        MutableObj<T> wrapper = container.get(key);
        if (null == wrapper || null == wrapper.get()) {
            container.remove(key);
            return null;
        }

        return wrapper.get();
    }

    /**
     * 根据一级原生key校验目标对象是否存在
     * @param key
     * @return
     */
    public boolean containsRawKey(RAW_KEY key) {
        return Optional.ofNullable(rawMap.get(key))
                .map(MutableObj::get)
                .map(Objects::nonNull)
                .orElse(false);
    }

    /**
     * 根据二级扩展key校验目标分桶是否存在。
     * 分桶不存在和分桶中元素为零均返回false
     *
     * @param extKey
     * @return
     */
    public boolean containsExtKey(EXT_KEY extKey) {
        Map<RAW_KEY, T> allByExtKey = this.getAllByExtKey(extKey);
        return (!CollUtil.isEmpty(allByExtKey));
    }

    /**
     * 根据一级拓原生key移除key与对象
     * @param key
     * @return
     */
    public T removeByRawKey(RAW_KEY key) {

        MutableObj<T> removed = rawMap.remove(key);

        if (null != removed) {
            removed.set(null);
        }

        return Optional.ofNullable(removed)
                .map(MutableObj::get)
                .orElse(null);
    }


    /**
     * 根据拓展key移除所有对应的对象
     * @param extKey
     * @return
     */
    public Map<RAW_KEY, T> removeAllByExtKey(EXT_KEY extKey) {
        ConcurrentHashMap<RAW_KEY, MutableObj<T>> removed = extMap.remove(extKey);

        HashMap<RAW_KEY, T> res = new HashMap<>();

        if (CollUtil.isEmpty(removed)) {
            if (null != removed) {
                this.extMap.remove(extKey);
            }

            return res;
        }

        for (Map.Entry<RAW_KEY, MutableObj<T>> entry : removed.entrySet()) {
            RAW_KEY rawKey = entry.getKey();
            MutableObj<T> value = entry.getValue();

            rawMap.remove(rawKey);

            if (null != value && null != value.get()) {
                res.put(rawKey, value.get());
            }
        }
        return res;
    }

    /**
     * 根据拓展key获取所有对应的对象
     * @param extKey
     * @return
     */
    public Map<RAW_KEY, T> getAllByExtKey(EXT_KEY extKey) {
        ConcurrentHashMap<RAW_KEY, MutableObj<T>> container = this.extMap.get(extKey);

        HashMap<RAW_KEY, T> res = new HashMap<>();

        if (CollUtil.isEmpty(container)) {
            if (null != container) {
                this.removeAllByExtKey(extKey);
            }

            return res;
        }

        for (Map.Entry<RAW_KEY, MutableObj<T>> entry : container.entrySet()) {
            RAW_KEY rawKey = entry.getKey();
            MutableObj<T> value = entry.getValue();

            if (null != value && null != value.get()) {
                res.put(rawKey, value.get());
            }

            if (null != value && null == value.get()) {
                container.remove(rawKey);
            }


        }

        return res;
    }



    public ExtendKeyBucketMap<RAW_KEY, EXT_KEY, T> getCopy() {
        return getCopy(false);
    }


    public ExtendKeyBucketMap<RAW_KEY, EXT_KEY, T> getCopy(boolean isConcurrent) {
        Map<EXT_KEY, ConcurrentHashMap<RAW_KEY, MutableObj<T>>> extMapNew = null;

        Map<RAW_KEY, MutableObj<T>> rawMapNew = null;

        if (isConcurrent) {
            extMapNew = new ConcurrentHashMap<>();
            rawMapNew = new ConcurrentHashMap<>();
        } else {
            extMapNew = new HashMap<>();
            rawMapNew = new HashMap<>();
        }

        // copy
        extMapNew.putAll(this.extMap);
        rawMapNew.putAll(this.rawMap);;

        return new ExtendKeyBucketMap<>(extMapNew, rawMapNew);
    }







}
